# -*- coding:utf-8 -*-
"""
@file name      : tensorboard_methods.py
@author         : QuZhang
@date           : 2020-12-23
@brief          : tensorboard方法使用(一) scalars and histogram
"""
import os
from tools.common_tools import set_seed
from torch.utils.tensorboard import SummaryWriter
import numpy as np
from matplotlib import pyplot as plt


os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
set_seed(1)

# ------------------- 0 SummaryWriter ------------
# flag = True
flag = False
if flag:
    log_dir = "./train_log/test_log_dir"
    # log_dir 设置后，comment无效
    # comment文件夹后缀，filename_suffix文件名后缀
    # writer = SummaryWriter(log_dir=log_dir, comment='_scalars', filename_suffix="12345678")
    writer = SummaryWriter(comment='_scalars', filename_suffix="123456789")

    for x in range(100):
        writer.add_scalar("y=pow_2_x", 2 ** x, x)
    writer.close()

# ---------------- 1 scalar and scalars ------------
# 正常函数曲线图
# flag = True
flag = False
if flag:
    log_dir = "./my_train_log/log_dir"
    max_epoch = 100
    writer = SummaryWriter(log_dir=log_dir, comment='test_comment', filename_suffix="test_suffix")

    for x in range(max_epoch):
        writer.add_scalar('y=2x', 2*x, x)
        writer.add_scalar("y=pow_2_x", 2**x, x)
        # 同时绘制多条曲线不同的曲线放在不同的文件夹
        # 第一个参数和字典的键值组成文件夹的名字
        writer.add_scalars('data/scalar_group', {"xsinx": x*np.sin(x),
                                                 'xcosx': x*np.cos(x)}, x)

    writer.close()

# ---------------- 2 histogram --------------
# 直方图，用于统计使用
flag = True
if flag:
    log_dir = "./my_train_log/log_dir"
    writer = SummaryWriter(log_dir=log_dir, comment='_histogram', filename_suffix='hist')

    for x in range(2):
        np.random.seed(x)
        data_union = np.arange(100)
        data_normal = np.random.normal(size=1000)

        writer.add_histogram('distribution union', data_union, x)
        writer.add_histogram('distribution normal', data_normal, x)

        plt.subplot(121).hist(data_union, label='union')
        plt.subplot(122).hist(data_normal, label='normal')
        plt.legend()
        plt.show()

    writer.close()

